import { useState, FormEvent, useEffect, useRef } from 'react';
import { flushSync } from 'react-dom';
import { useSWRWithAuth as useSWR } from 'renderer/lib/authContext';

import * as chatAPI from 'renderer/lib/transformerlab-api-sdk';
import Markdown from 'react-markdown';
import remarkGfm from 'remark-gfm';
import {
  Button,
  FormControl,
  FormHelperText,
  FormLabel,
  Input,
  Modal,
  ModalDialog,
  Stack,
  Tab,
  TabList,
  TabPanel,
  Tabs,
  Sheet,
  Switch,
} from '@mui/joy';
import DynamicPluginForm from '../DynamicPluginForm';
import TrainingModalDataTab from './TraningModalDataTab';

import AvailableFieldsImage from 'renderer/img/show-available-fields.png';

import { generateFriendlyName } from 'renderer/lib/utils';
import OneTimePopup from 'renderer/components/Shared/OneTimePopup';
import TrainingModalDataTemplatingTab from './TrainingModalDataTemplatingTab';
import SafeJSONParse from 'renderer/components/Shared/SafeJSONParse';
import { fetcher } from 'renderer/lib/transformerlab-api-sdk';

function isValidChatMessageList(input: any): boolean {
  try {
    const parsed = typeof input === 'string' ? JSON.parse(input) : input;

    if (!Array.isArray(parsed)) return false;

    for (const msg of parsed) {
      if (
        typeof msg !== 'object' ||
        msg === null ||
        typeof msg.role !== 'string' ||
        typeof msg.content !== 'string'
      ) {
        return false;
      }
    }

    return true;
  } catch (e) {
    return false;
  }
}

function PluginIntroduction({ experimentInfo, pluginId }) {
  const { data, error, isLoading } = useSWR(
    chatAPI.Endpoints.Experiment.ScriptGetFile(
      experimentInfo?.id,
      pluginId,
      'info.md',
    ),
    fetcher,
  );

  return (
    <>
      <Markdown remarkPlugins={[remarkGfm]} className="editableSheetContent">
        {data && data != 'FILE NOT FOUND'
          ? data
          : 'No description for this plugin is available.'}
      </Markdown>
    </>
  );
}

export default function TrainingModalLoRA({
  open,
  onClose,
  experimentInfo,
  task_id,
  pluginId,
}: {
  open: boolean;
  onClose: () => void;
  experimentInfo: any;
  task_id?: string;
  pluginId: string;
}) {
  // Store the current selected Dataset in this modal
  const [selectedDataset, setSelectedDataset] = useState(null);
  const [config, setConfig] = useState({});
  const [nameInput, setNameInput] = useState('');
  const [currentTab, setCurrentTab] = useState(0);
  const [sweepConfig, setSweepConfig] = useState<{ [key: string]: string[] }>(
    {},
  );
  const [isRunSweeps, setIsRunSweeps] = useState(false);

  const [applyChatTemplate, setApplyChatTemplate] = useState(false);
  const [chatColumn, setChatColumn] = useState('');
  const [formattingTemplate, setFormattingTemplate] = useState('');
  const [formattingChatTemplate, setFormattingChatTemplate] = useState('');
  const nameInputRef = useRef<HTMLInputElement | null>(null);

  // Preserve cursor position in controlled input
  const handleNameChange = (e: React.ChangeEvent<HTMLInputElement>) => {
    // Before updating state, capture cursor position & scroll so we can restore them.
    // React re-renders controlled inputs, which normally resets the cursor (causing jumps).

    const inputEl =
      (nameInputRef.current as HTMLInputElement | null) ||
      (e.target as HTMLInputElement);

    const start = inputEl?.selectionStart ?? 0;
    const end = inputEl?.selectionEnd ?? start;
    const scrollLeft = inputEl?.scrollLeft ?? 0;
    const value = (e.target as HTMLInputElement).value;

    flushSync(() => {
      setNameInput(value); // update input value immediately
    });

    // After the re-render, restore cursor & scroll so typing feels natural in controlled input.
    const cur = nameInputRef.current;
    if (cur && document.activeElement === cur) {
      try {
        const max = cur.value.length;
        cur.setSelectionRange(Math.min(start, max), Math.min(end, max));
        cur.scrollLeft = scrollLeft;
      } catch {
        // ignore if browser doesn't allow programmatic cursor updates
      }
    }
  };

  // Fetch training type with useSWR
  const { data: trainingTypeData } = useSWR(
    experimentInfo?.id
      ? chatAPI.Endpoints.Experiment.ScriptGetFile(
          experimentInfo.id,
          pluginId,
          'index.json',
        )
      : null,
    fetcher,
  );

  let trainingType = 'LoRA';
  if (
    trainingTypeData &&
    trainingTypeData !== 'undefined' &&
    trainingTypeData !== 'FILE NOT FOUND' &&
    trainingTypeData.length > 0
  ) {
    trainingType = SafeJSONParse(trainingTypeData)?.train_type || 'LoRA';
  }

  let runSweeps = false;
  if (
    trainingTypeData &&
    trainingTypeData !== 'undefined' &&
    trainingTypeData !== 'FILE NOT FOUND' &&
    trainingTypeData.length > 0
  ) {
    const parsedData = SafeJSONParse(trainingTypeData, {});
    if (Array.isArray(parsedData?.supports)) {
      runSweeps = parsedData.supports.includes('sweeps');
    }
  }

  // Fetch available datasets from the API
  const {
    data: datasets,
    error: datasetsError,
    isLoading: datasetsIsLoading,
  } = useSWR(chatAPI.Endpoints.Dataset.LocalList(), fetcher);

  const {
    data: templateData,
    error: templateError,
    isLoading: templateIsLoading,
    mutate: templateMutate,
  } = useSWR(
    task_id ? chatAPI.Endpoints.Tasks.GetByID(task_id) : null,
    fetcher,
  );

  async function updateTask(
    task_id: string,
    name: string,
    inputs: string,
    config: string,
    outputs: string,
  ) {
    const configBody = {
      name: name,
      inputs: inputs,
      config: config,
      outputs: outputs,
    };
    const response = await chatAPI.authenticatedFetch(
      chatAPI.Endpoints.Tasks.UpdateTask(task_id),
      {
        method: 'PUT',
        headers: {
          'Content-Type': 'application/json',
          accept: 'application/json',
        },
        body: JSON.stringify(configBody),
      },
    );
    const result = await response.json();
    return result;
  }

  async function createNewTask(
    name: string,
    plugin: string,
    experimentId: string,
    inputs: string,
    config: string,
    outputs: string,
  ) {
    const configBody = {
      name: name,
      plugin: plugin,
      experiment_id: experimentId,
      inputs: inputs,
      config: config,
      outputs: outputs,
      type: 'TRAIN',
    };
    const response = await chatAPI.authenticatedFetch(
      chatAPI.Endpoints.Tasks.NewTask(),
      {
        method: 'PUT',
        headers: {
          'Content-Type': 'application/json',
          accept: 'application/json',
        },
        body: JSON.stringify(configBody),
      },
    );
    const result = await response.json();
    return result;
  }
  // Handle modal open/close state - reset form when modal closes
  useEffect(() => {
    if (open) {
      if (!task_id || task_id === '') {
        setNameInput(generateFriendlyName());
      } else {
        // If we have a task_id and modal is opening, refetch the template data
        templateMutate();
      }
    } else {
      // Reset all form state when modal closes
      setSelectedDataset(null);
      setConfig({});
      setNameInput('');
      setCurrentTab(0);
      setSweepConfig({});
      setIsRunSweeps(false);
      setApplyChatTemplate(false);
      setChatColumn('');
      setFormattingTemplate('');
      setFormattingChatTemplate('');
    }
  }, [open, task_id, templateMutate]);

  // Whenever template data updates, we need to update state variables used in the form.
  useEffect(() => {
    if (!templateData) {
      setSelectedDataset(null);
      setConfig({});
      setNameInput(generateFriendlyName());
      setFormattingTemplate('');
      setFormattingChatTemplate('');
      setApplyChatTemplate(false);
      setChatColumn('');
      setSweepConfig({});
      setIsRunSweeps(false);
      return;
    }

    // Parse inputs to get dataset_name
    const parsedInputs = SafeJSONParse(templateData.inputs, {});

    // Parse config
    let parsedConfig = null;
    if (templateData.config) {
      if (typeof templateData.config === 'string') {
        parsedConfig = SafeJSONParse(templateData.config, {});
      } else if (typeof templateData.config === 'object') {
        parsedConfig = templateData.config;
      }
    }

    if (parsedConfig) {
      // Dataset name can be in either config or inputs
      const datasetName =
        parsedConfig.dataset_name || parsedInputs.dataset_name;
      setSelectedDataset(datasetName || null);
      setConfig(parsedConfig);
      setNameInput(templateData.name || generateFriendlyName());

      // Set template fields
      if (parsedConfig.formatting_template) {
        setFormattingTemplate(parsedConfig.formatting_template);
      }
      if (parsedConfig.formatting_chat_template) {
        setFormattingChatTemplate(parsedConfig.formatting_chat_template);
      }
      if (parsedConfig.apply_chat_template !== undefined) {
        setApplyChatTemplate(parsedConfig.apply_chat_template);
      }
      if (parsedConfig.chat_column) {
        setChatColumn(parsedConfig.chat_column);
      }

      if (parsedConfig.sweep_config) {
        setSweepConfig(SafeJSONParse(parsedConfig.sweep_config, {}));
      } else {
        setSweepConfig({});
      }
      if (parsedConfig.run_sweeps) {
        setIsRunSweeps(parsedConfig.run_sweeps);
      } else {
        setIsRunSweeps(false);
      }
    } else {
      // This case is for when we are creating a new template or config is missing
      setSelectedDataset(null);
      setConfig({});
      setNameInput(templateData?.name || generateFriendlyName());
      setFormattingTemplate('');
      setFormattingChatTemplate('');
      setApplyChatTemplate(false);
      setChatColumn('');
      setSweepConfig({});
      setIsRunSweeps(false);
    }
  }, [templateData]);
  // Once you have a dataset selected, we use SWR's dependency mode to fetch the
  // Dataset's info. Note how useSWR is declared as a function -- this is is how
  // the dependency works. If selectedDataset errors, the fetcher knows to not run.
  const { data: currentDatasetInfo, isLoading: currentDatasetInfoIsLoading } =
    useSWR(() => {
      if (selectedDataset === null) {
        return null;
      }
      return chatAPI.Endpoints.Dataset.Info(selectedDataset);
    }, fetcher);
  const currentModel = experimentInfo?.config?.foundation_filename
    ? experimentInfo?.config?.foundation_filename
    : experimentInfo?.config?.foundation;

  function injectIntoTemplate(key) {
    // Add the key to the textbox with id "template"
    const template =
      document.getElementById('training-form')?.elements['formatting_template'];

    if (template === undefined) return;

    const cursorPosition = template.selectionStart;
    const templateText = template.value;
    const newText =
      templateText.slice(0, cursorPosition) +
      `{{${key}}}` +
      templateText.slice(cursorPosition);
    template.value = newText;
  }

  if (!experimentInfo?.id) {
    return 'Select an Experiment';
  }

  // eslint-disable-next-line react/no-unstable-nested-components
  function TrainingModalFirstTab() {
    return (
      <Stack spacing={2}>
        <FormControl>
          <FormLabel>Training Template Name</FormLabel>
          <Input
            required
            autoFocus
            value={nameInput} //Value needs to be stored in a state variable otherwise it will not update on change/update
            onChange={handleNameChange}
            slotProps={{ input: { ref: nameInputRef } }}
            name="template_name"
            size="lg"
          />
          <FormHelperText>
            Give this specific training recipe a unique name
          </FormHelperText>
        </FormControl>
        <FormLabel>Info</FormLabel>
        <Sheet color="neutral" variant="soft">
          <Stack direction="column" justifyContent="space-evenly" gap={2} p={2}>
            <FormControl sx={{ flex: 1 }}>
              <FormLabel>Base Model:</FormLabel>
              <Input readOnly value={currentModel} variant="soft" />
            </FormControl>
            <FormControl sx={{ flex: 1 }}>
              <FormLabel>Architecture:</FormLabel>
              <Input
                readOnly
                value={experimentInfo?.config?.foundation_model_architecture}
                variant="soft"
              />
            </FormControl>
            <FormControl sx={{ flex: 1 }}>
              <FormLabel>Plugin:</FormLabel>
              <Input
                readOnly
                value={pluginId}
                variant="soft"
                name="plugin_name"
              />
            </FormControl>
            <input hidden value={currentModel} name="model_name" readOnly />
            <input
              hidden
              value={experimentInfo?.config?.foundation_model_architecture}
              name="model_architecture"
              readOnly
            />
            <input
              hidden
              value={experimentInfo?.config?.foundation_filename}
              name="foundation_model_file_path"
              readOnly
            />
            <input
              hidden
              value={experimentInfo?.config?.embedding_model}
              name="embedding_model"
              readOnly
            />
            <input
              hidden
              value={experimentInfo?.config?.embedding_model_architecture}
              name="embedding_model_architecture"
              readOnly
            />
            <input
              hidden
              value={experimentInfo?.config?.embedding_model_filename}
              name="embedding_model_file_path"
              readOnly
            />
          </Stack>
        </Sheet>
      </Stack>
    );
  }
  // eslint-disable-next-line react/no-unstable-nested-components
  function SweepConfigTab({ trainingTypeData, sweepConfig, setSweepConfig }) {
    const [newParam, setNewParam] = useState('');
    const [newValues, setNewValues] = useState('');

    // Parse parameters from trainingTypeData
    let availableParameters = [];
    let parameterTypes = {};
    try {
      if (trainingTypeData && trainingTypeData !== 'undefined') {
        const parsedData = SafeJSONParse(trainingTypeData, {});
        // Extract parameter names and their types from the training type data
        if (parsedData?.parameters) {
          availableParameters = Object.keys(parsedData.parameters);
          // Store parameter types for conversion later
          Object.entries(parsedData.parameters).forEach(([param, config]) => {
            if (config && typeof config === 'object' && 'type' in config) {
              parameterTypes[param] = config.type;
            }
          });
        }
      }
    } catch (error) {
      console.error('Error parsing training type data:', error);
    }

    // Filter out parameters that are already in sweepConfig
    const unusedParameters = availableParameters.filter(
      (param) => !Object.keys(sweepConfig).includes(param),
    );

    const addSweepParam = () => {
      if (newParam && newValues.trim()) {
        // Split values by comma and trim whitespace
        const valuesArray = newValues.split(',').map((val) => {
          const trimmedValue = val.trim();

          // Convert to number if parameter type is number or integer
          const paramType = parameterTypes[newParam];
          if (paramType === 'number' || paramType === 'integer') {
            const numValue = Number(trimmedValue);
            return isNaN(numValue) ? trimmedValue : numValue;
          }

          return trimmedValue;
        });

        setSweepConfig((prev) => ({
          ...prev,
          [newParam]: valuesArray,
        }));

        // Reset input fields
        setNewParam('');
        setNewValues('');
      }
    };

    const removeParam = (paramToRemove) => {
      setSweepConfig((prev) => {
        const updated = { ...prev };
        delete updated[paramToRemove];
        return updated;
      });
    };

    return (
      <Stack spacing={3}>
        <Sheet sx={{ p: 2, borderRadius: 'sm' }} variant="outlined">
          <Stack spacing={2}>
            <Stack
              direction="row"
              justifyContent="space-between"
              alignItems="center"
            >
              <FormLabel>Run Hyperparameter Sweeps</FormLabel>
              <Switch
                checked={isRunSweeps}
                onChange={(event) => setIsRunSweeps(event.target.checked)}
                color={isRunSweeps ? 'success' : 'neutral'}
              />
            </Stack>
            <FormHelperText>
              Enable this to perform hyperparameter sweeps using the parameters
              defined below.
            </FormHelperText>
          </Stack>
        </Sheet>
        <Sheet sx={{ p: 2, borderRadius: 'sm' }} variant="outlined">
          <Stack spacing={2}>
            <FormLabel>Add Parameter Sweep</FormLabel>
            <FormHelperText>
              Define parameters to sweep during training. Each parameter can
              have multiple values to try. Selecting a hyperparameter will
              override the values set for it in the Plugin Config tab.
            </FormHelperText>

            <Stack direction="row" spacing={2} alignItems="flex-start">
              <FormControl sx={{ minWidth: '200px' }}>
                <FormLabel>Parameter</FormLabel>
                <select
                  value={newParam}
                  onChange={(e) => setNewParam(e.target.value)}
                  style={{
                    padding: '10px',
                    borderRadius: '6px',
                    border: '1px solid #ccc',
                    width: '100%',
                    backgroundColor: 'transparent',
                  }}
                >
                  <option value="">Select a parameter</option>
                  {unusedParameters.map((param) => (
                    <option key={param} value={param}>
                      {param}
                    </option>
                  ))}
                </select>
              </FormControl>

              <FormControl sx={{ flex: 1 }}>
                <FormLabel>Sweep Values (comma separated)</FormLabel>
                <Input
                  value={newValues}
                  onChange={(e) => setNewValues(e.target.value)}
                  placeholder="e.g. 2,4,8,16"
                />
                <FormHelperText>
                  Enter values separated by commas
                </FormHelperText>
              </FormControl>

              <Button
                sx={{ mt: 3 }}
                onClick={addSweepParam}
                disabled={!newParam || !newValues.trim()}
              >
                Add Parameter
              </Button>
            </Stack>
          </Stack>
        </Sheet>

        {Object.keys(sweepConfig).length > 0 && (
          <Sheet sx={{ p: 2, borderRadius: 'sm' }} variant="outlined">
            <FormLabel>Current Sweep Configuration</FormLabel>
            <Stack spacing={2} mt={1}>
              {Object.entries(sweepConfig).map(([param, values]) => (
                <Sheet
                  key={param}
                  sx={{ p: 2, borderRadius: 'sm' }}
                  variant="soft"
                >
                  <Stack
                    direction="row"
                    justifyContent="space-between"
                    alignItems="center"
                  >
                    <Stack>
                      <FormLabel>{param}</FormLabel>
                      <FormHelperText>
                        Values: {values.join(', ')}
                      </FormHelperText>
                    </Stack>
                    <Button
                      color="danger"
                      variant="soft"
                      size="sm"
                      onClick={() => removeParam(param)}
                    >
                      Remove
                    </Button>
                  </Stack>
                </Sheet>
              ))}
            </Stack>
          </Sheet>
        )}

        <input
          type="hidden"
          name="sweep_config"
          value={JSON.stringify(sweepConfig)}
        />
        <input type="hidden" name="run_sweeps" value={isRunSweeps.toString()} />
      </Stack>
    );
  }

  return (
    <Modal open={open}>
      <ModalDialog
        sx={{
          width: '95dvw',
          transform: 'translateX(-50%)', // This undoes the default translateY that centers vertically
          top: '5dvh',
          overflow: 'auto',
          maxHeight: '92dvh',
          minHeight: '70dvh',
          height: '100%',
        }}
      >
        <form
          id="training-form"
          style={{
            display: 'flex',
            flexDirection: 'column',
            height: '100%',
            justifyContent: 'space-between',
          }}
          noValidate
          onSubmit={async (event: FormEvent<HTMLFormElement>) => {
            event.preventDefault();

            const formEl = event.currentTarget;

            // If the form is not valid
            if (!formEl.checkValidity()) {
              // Get the first invalid element
              const firstInvalid = formEl.querySelector(
                ':invalid',
              ) as HTMLElement | null;

              if (firstInvalid) {
                // Get the tab containing the invalid element
                const tabPanel = firstInvalid.closest(
                  '[data-tab-index]',
                ) as HTMLElement | null;

                if (tabPanel) {
                  const idxAttr = tabPanel.getAttribute('data-tab-index');
                  // Get the index of the tab
                  const idx = idxAttr ? Number(idxAttr) : NaN;

                  if (!Number.isNaN(idx)) {
                    // Switch to the tab
                    flushSync(() => setCurrentTab(idx));
                  }
                }
              }

              requestAnimationFrame(() => {
                // Trigger browser validation
                formEl.reportValidity();

                // Get invalid element again (after triggering browser validation)
                // reportValidity() can cause the browser to update validation state
                const invalidEl = formEl.querySelector(
                  ':invalid',
                ) as HTMLElement | null;

                if (invalidEl) {
                  // Scroll to the element and focus it
                  invalidEl.scrollIntoView({
                    behavior: 'smooth',
                    block: 'center',
                  });
                  invalidEl.focus();
                }
              });

              return;
            }
            const formData = new FormData(event.currentTarget);
            let formJson = Object.fromEntries((formData as any).entries());

            // Merge data from DynamicPluginForm (stored in window.currentFormData)
            if ((window as any).currentFormData) {
              Object.assign(formJson, (window as any).currentFormData);
            }

            // Ensure nameInput state is captured as template_name
            formJson.template_name = nameInput;

            // Include selectedDataset if it exists
            if (selectedDataset) {
              formJson.dataset_name = selectedDataset;
            }

            // Include chat template settings
            if (applyChatTemplate) {
              formJson.apply_chat_template = applyChatTemplate;
            }
            if (chatColumn) {
              formJson.chat_column = chatColumn;
            }

            // Include template values
            if (formattingTemplate) {
              formJson.formatting_template = formattingTemplate;
            }
            if (formattingChatTemplate) {
              formJson.formatting_chat_template = formattingChatTemplate;
            }

            formJson.type = trainingType;
            formJson.chatml_formatted_column = chatColumn;
            // Add sweep config to form data
            if (Object.keys(sweepConfig).length > 0) {
              formJson.sweep_config = JSON.stringify(sweepConfig);
            }
            if (formJson.run_sweeps) {
              formJson.run_sweeps = formJson.run_sweeps === 'true';
            } else {
              formJson.run_sweeps = false;
            }

            if (templateData && task_id) {
              //Only update if we are currently editing a template
              // For all keys in templateData.inputs that are in formJson, set the value from formJson
              const templateDataInputs = SafeJSONParse(templateData.inputs, {});
              const templateDataOutputs = SafeJSONParse(
                templateData.outputs,
                {},
              );
              for (const key in templateDataInputs) {
                if (
                  key in formJson &&
                  templateDataInputs[key] != formJson[key]
                ) {
                  templateDataInputs[key] = formJson[key];
                }
              }
              // For all keys in templateData.outputs that are in formJson, set the value from formJson
              for (const key in templateDataOutputs) {
                if (
                  key in formJson &&
                  templateDataOutputs[key] != formJson[key]
                ) {
                  templateDataOutputs[key] = formJson[key];
                }
              }

              await updateTask(
                task_id,
                formJson.template_name,
                JSON.stringify(templateDataInputs),
                JSON.stringify(formJson),
                JSON.stringify(templateDataOutputs),
              );
              templateMutate(); //Need to mutate template data after updating
            } else {
              await createNewTask(
                formJson.template_name,
                formJson.plugin_name,
                experimentInfo?.id,
                JSON.stringify({
                  model_name: formJson.model_name,
                  model_architecture: formJson.model_architecture,
                  dataset_name: formJson.dataset_name,
                }),
                JSON.stringify(formJson),
                JSON.stringify({
                  adaptor_name: formJson.adaptor_name,
                }),
              );
            }
            onClose();
          }}
        >
          <Tabs
            aria-label="Training Template Tabs"
            value={currentTab}
            onChange={(event, newValue) => setCurrentTab(newValue)}
            sx={{ borderRadius: 'lg', display: 'flex', overflow: 'hidden' }}
          >
            <TabList>
              <Tab>Introduction</Tab>
              <Tab>Name</Tab>
              <Tab>Dataset</Tab>
              <Tab>Data Template</Tab>
              <Tab>Plugin Config</Tab>
              {runSweeps && <Tab>Sweep Config</Tab>}
            </TabList>
            <TabPanel
              data-tab-index={0}
              value={0}
              sx={{ p: 2, overflow: 'auto' }}
            >
              <PluginIntroduction
                experimentInfo={experimentInfo}
                pluginId={pluginId}
              />
            </TabPanel>
            <TabPanel
              data-tab-index={1}
              value={1}
              sx={{ p: 2, overflow: 'auto' }}
              keepMounted
            >
              <TrainingModalFirstTab />
            </TabPanel>
            <TabPanel
              data-tab-index={3}
              value={3}
              sx={{ p: 2, overflow: 'auto' }}
              keepMounted
            >
              <TrainingModalDataTemplatingTab
                selectedDataset={selectedDataset}
                currentDatasetInfo={currentDatasetInfo}
                templateData={templateData}
                injectIntoTemplate={injectIntoTemplate}
                experimentInfo={experimentInfo}
                pluginId={pluginId}
                applyChatTemplate={applyChatTemplate}
                setApplyChatTemplate={setApplyChatTemplate}
                chatColumn={chatColumn}
                setChatColumn={setChatColumn}
                formattingTemplate={formattingTemplate}
                setFormattingTemplate={setFormattingTemplate}
                formattingChatTemplate={formattingChatTemplate}
                setFormattingChatTemplate={setFormattingChatTemplate}
              />
            </TabPanel>
            <TabPanel
              data-tab-index={2}
              value={2}
              sx={{ p: 2, overflow: 'auto' }}
              keepMounted
            >
              <>
                {currentTab === 2 && (
                  <OneTimePopup title="How to Create a Training Template:">
                    Use the <b>Available Fields</b> to populate the template
                    fields on this screen. For each template field, you can type
                    any text, and when you want to inject text from your
                    dataset, add the field name, surrounded with curly braces
                    like this: &#123;&#123; example &#125;&#125; .
                    <br />
                    <br />
                    <img
                      src={AvailableFieldsImage}
                      alt="Available Fields"
                      width="400"
                    />
                    <br />
                    <br />
                    The Avaiable Fields will change dynamically based on the
                    columns in your selected dataset.
                  </OneTimePopup>
                )}

                <TrainingModalDataTab
                  datasetsIsLoading={datasetsIsLoading}
                  datasets={datasets}
                  selectedDataset={selectedDataset}
                  setSelectedDataset={setSelectedDataset}
                  currentDatasetInfoIsLoading={currentDatasetInfoIsLoading}
                  currentDatasetInfo={currentDatasetInfo}
                  templateData={templateData}
                  injectIntoTemplate={injectIntoTemplate}
                  experimentInfo={experimentInfo}
                  pluginId={pluginId}
                  displayMessage=""
                />
              </>
            </TabPanel>
            <TabPanel
              data-tab-index={4}
              value={4}
              sx={{ p: 2, overflow: 'auto' }}
              keepMounted
            >
              <DynamicPluginForm
                experimentInfo={experimentInfo}
                plugin={pluginId}
                config={config}
              />
            </TabPanel>
            {runSweeps && (
              <TabPanel
                data-tab-index={5}
                value={5}
                sx={{ p: 2, overflow: 'auto' }}
                keepMounted
              >
                <SweepConfigTab
                  trainingTypeData={trainingTypeData}
                  sweepConfig={sweepConfig}
                  setSweepConfig={setSweepConfig}
                />
              </TabPanel>
            )}
          </Tabs>
          <Stack spacing={2} direction="row" justifyContent="flex-end">
            <Button color="danger" variant="soft" onClick={() => onClose()}>
              Cancel
            </Button>
            <Button
              variant="soft"
              type="submit"
              color="success"
              disabled={
                applyChatTemplate &&
                (!chatColumn ||
                  chatColumn.trim() === '' ||
                  !isValidChatMessageList(chatColumn))
              }
            >
              Save Training Template
            </Button>
          </Stack>
        </form>
        {/* {JSON.stringify(config, null, 2)} */}
      </ModalDialog>
    </Modal>
  );
}
